#!/usr/bin/env python

import gradio as gr
from PIL import Image
import os
import requests
import json
import argparse
import torch
import base64
import time
from model import is_chinese, generate_input, chat
from transformers import AutoTokenizer
from finetune_XrayGLM import FineTuneVisualGLMModel
from sat.model import AutoModel
from sat.model.mixins import CachedAutoregressiveMixin
from sat.quantization.kernels import quantize
from datetime import datetime  

# 全局变量
global model, tokenizer, current_report, patient_info

# 通义千问API配置（请确认以下信息正确）
QWEN_API_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1"
QWEN_API_KEY = "sk-728bc98223eb4321babd2618503eae3a"  
HEADERS = {
    "Content-Type": "application/json",
    "Authorization": f"Bearer {QWEN_API_KEY}",
    "X-DashScope-Async": "enable"
}

# 初始化患者信息
patient_info = {
    "name": "未填写",
    "age": "未知",
    "gender": "未填写",
    "history": "无"
}


def preprocess_medical_image(image, target_size=512):
    """医学影像预处理"""
    if hasattr(image, '_getexif'):
        exif = image._getexif()
        if exif:
            orientation = exif.get(0x0112, 1)
            if orientation == 3:
                image = image.rotate(180, expand=True)
            elif orientation == 6:
                image = image.rotate(270, expand=True)
            elif orientation == 8:
                image = image.rotate(90, expand=True)
    
    width, height = image.size
    scale = target_size / max(width, height)
    new_size = (int(width*scale), int(height*scale))
    image = image.resize(new_size, Image.LANCZOS)
    
    if image.mode != 'RGB':
        image = image.convert('RGB')
    
    return image

def call_qwen_api(image_path, prompt):
    """调用通义千问多模态API（同步模式）"""
    try:
        with open(image_path, "rb") as image_file:
            image_data = base64.b64encode(image_file.read()).decode('utf-8')
        
        payload = {
            "model": "qwen-vl-max-latest",
            "input": {
                "messages": [
                    {
                        "role": "user",
                        "content": [
                            {"image": f"data:image/jpeg;base64,{image_data}"},
                            {"text": prompt}
                        ]
                    }
                ]
            },
            "parameters": {
                "temperature": 0.5,
                "top_p": 0.9,
                "max_tokens": 3000
            }
        }
        
        # 修改为同步API端点
        sync_api_url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation"
        
        # 移除异步请求头
        headers = {
            "Authorization": f"Bearer {QWEN_API_KEY}",
            "Content-Type": "application/json"
        }
        
        response = requests.post(sync_api_url, headers=headers, json=payload)
        
        if response.status_code == 200:
            return response.json()["output"]["choices"][0]["message"]["content"]
        else:
            error_msg = f"""
            API请求失败！
            状态码：{response.status_code}
            错误信息：{response.text}
            可能原因：
            1. API密钥无效或过期
            2. 账户未开通多模态服务
            3. 请求参数不符合要求
            4. 账号余额不足
            """
            return error_msg

    except Exception as e:
        return f"API调用异常：{str(e)}"




def generate_text_with_image(input_text, image, history=[], request_data=dict(), is_zh=True):
    try:
        input_para = {
            "max_length": 2048,
            "min_length": 50,
            "temperature": 0.8,
            "top_p": 0.4,
            "top_k": 100,
            "repetition_penalty": 1.2
        }
        input_para.update(request_data)

        input_data = generate_input(input_text, image, history, input_para, image_is_encoded=False)
        input_image, gen_kwargs = input_data['input_image'], input_data['gen_kwargs']
        
        with torch.no_grad():
            answer, history, _ = chat(None, model, tokenizer, input_text, 
                                     history=history, image=input_image,
                                     max_length=gen_kwargs['max_length'],
                                     top_p=gen_kwargs['top_p'],
                                     top_k=gen_kwargs['top_k'],
                                     temperature=gen_kwargs['temperature'],
                                     english=not is_zh)
        return answer
    except Exception as e:
        return str(e)

# def generate_medical_report(history):
#     """生成结构化诊疗报告"""
#     global patient_info
#     report = f"# 医学影像诊疗报告\n\n"
#     report += f"**患者姓名**: {patient_info['name']}  "
#     report += f"**年龄**: {patient_info['age']}  "
#     report += f"**性别**: {patient_info['gender']}\n\n"
#     report += f"**病史摘要**: {patient_info['history']}\n\n"

#     # 提取关键信息
#     diagnosis = next((msg[1] for msg in history if "初步诊断" in msg[0]), "")
#     findings = []
#     recommendations = []
    
#     for q, a in history:
#         if "影像表现" in a:
#             findings.extend(a.split("影像表现：")[-1].split("\n"))
#         if "建议" in a:
#             recommendations.extend(a.split("建议：")[-1].split("\n"))
    
#     # 报告结构
#     report += "## 影像学表现\n"
#     report += "\n".join([f"- {f.strip()}" for f in findings if f.strip()]) or "无显著异常表现\n"
    
#     report += "\n\n## 初步诊断\n"
#     report += "\n".join([f"- {d.strip()}" for d in diagnosis.split("\n") if d.strip()]) + "\n"
    
#     report += "\n\n## 处理建议\n"
#     report += "\n".join([f"- {r.strip()}" for r in recommendations if r.strip()]) or "建议结合临床进一步检查\n"
    
#     report += "\n\n## 综合评估\n"
#     report += "基于影像学表现和AI分析，建议临床医生结合患者病史和其他检查结果进行综合判断。\n\n"
#     report += "**报告生成时间**: " + gr.utils.get_datetime() + "\n"
    
#     return report

import re

def sanitize_response(response):
    """增强版响应清洗和格式优化"""
    # 提取文本内容
    if isinstance(response, dict) and 'text' in response:
        text = response['text']
    else:
        text = str(response)
    
    # 清理特殊字符和格式符号
    text = re.sub(r'\*{5,}', '', text)  # 清除连续5个以上的星号
    text = re.sub(r'\[\{\{\'text\': \'###.*?\}\}', '', text)  
    text = re.sub(r'\'\}\}\]', '', text)
    text = re.sub(r'\bv[nm]\b', '', text)  # 清除vn/vm等无效后缀
    text = re.sub(r'\u3000', ' ', text)  # 替换全角空格
    text = text.replace('\\n', '\n')  # 转换转义换行符
    
    # 结构化格式优化
    text = re.sub(r'【(.*?)】', r'\n## \1\n', text)  # 转换章节标题
    text = re.sub(r'\*{3,}(\d+\.)', r'\n\1', text)  # 清理条目编号前的星号
    text = re.sub(r'(\d+\.)([^\n])', r'\1 \2', text)  # 确保编号后带空格
    
    # 转义特殊字符
    text = text.replace("{", "{{").replace("}", "}}")
    text = text.replace("data:image/", "data-image-")
    
    # 智能分段处理
    text = '\n'.join([line.strip() for line in text.split('\n') if line.strip()])
    
    return text[:2500]  # 适当增加长度限制

def format_diagnosis(text):
    """二次结构化处理"""
    # 优化列表格式
    text = re.sub(r'(\d+\.\s*)([^\n]+)', r'- \2\n', text)
    # 添加空行分隔
    text = re.sub(r'\n##', r'\n\n##', text)
    # 清理多余换行
    text = re.sub(r'\n{3,}', r'\n\n', text)
    return text

def request_model(input_text, temperature, top_p, image_prompt, result_previous):
    global current_report, patient_info
    
    result_text = [(ele[0], ele[1]) for ele in result_previous]
    result_text = [ele for ele in result_text if ele[0] or ele[1]]
    
    try:
        if image_prompt and not result_previous:
            diagnosis_prompt = """作为资深放射科医生，请严格按以下结构化格式分析：

                                    ## 影像表现
                                    1. 主要异常：
                                    - 部位：[明确解剖位置]
                                    - 形态：[具体形态描述]
                                    - 密度/信号特征：[量化指标]
                                    2. 重要阴性发现
                                    [列出重要阴性发现]
                                    3. 对比既往变化
                                    [对比结论/无法对比说明]

                                    ## 初步诊断
                                    1. 最可能诊断（按概率排序）
                                    [按概率降序排列，每个诊断单独一行]
                                    2. 鉴别诊断及排除依据
                                    [列出需排除的疾病及依据]

                                    ## 建议
                                    1. 推荐影像检查
                                    [具体检查项目]
                                    2. 实验室检查
                                    [检验项目]
                                    3. 临床处理方案
                                    [治疗方案]
                                    4. 随访计划
                                    [复查时间及方式]"""
                                                
            diagnosis = call_qwen_api(image_prompt, diagnosis_prompt)
            clean_diagnosis = format_diagnosis(sanitize_response(diagnosis))
            result_text.append(("AI初步诊断", clean_diagnosis))
            current_report = clean_diagnosis
            
        if image_prompt:
            image = Image.open(image_prompt)
            request_para = {"temperature": temperature, "top_p": top_p}
            answer = generate_text_with_image(input_text, image, result_text.copy(), request_para, is_chinese(input_text))
            formatted_answer = format_diagnosis(sanitize_response(answer))
            result_text.append((input_text, formatted_answer))
            
        return "", result_text
    except Exception as e:
        error_msg = f"系统错误：{type(e).__name__}\n详情：{str(e)}"
        return "", result_text + [(input_text, error_msg)]

def generate_medical_report(history):
    """根据对话历史生成结构化报告"""
    global patient_info
    try:
        report = "# 医学影像诊疗报告\n\n"
        
        # ================= 患者信息 =================
        report += "## 患者基本信息\n"
        report += f"- **姓名**: {patient_info['name']}\n"
        report += f"- **年龄**: {patient_info['age']}\n"
        report += f"- **性别**: {patient_info['gender']}\n"
        report += f"- **病史**: {patient_info['history']}\n\n"

        # ================= 解析对话历史 =================
        diagnosis_content = "\n".join([msg[1] for msg in history if msg[0] == "AI初步诊断"])
        
        # 初始化结构化数据容器
        structured_data = {
            "影像表现": {"主要异常": {}, "阴性发现": [], "对比变化": ""},
            "初步诊断": {"最可能诊断": [], "鉴别诊断": []},
            "处理建议": {"影像检查": [], "实验室检查": [], "临床处理": [], "随访计划": []}
        }

        # 使用状态机进行内容解析
        current_section = None
        current_subsection = None
        for line in diagnosis_content.split('\n'):
            line = line.strip()
            
            # 检测二级标题
            if line.startswith("## 影像表现"):
                current_section = "影像表现"
            elif line.startswith("## 初步诊断"):
                current_section = "初步诊断"
            elif line.startswith("## 建议"):
                current_section = "处理建议"
            
            # 影像表现解析
            elif current_section == "影像表现":
                if line.startswith("1. 主要异常："):
                    current_subsection = "主要异常"
                elif line.startswith("- 部位："):
                    structured_data["影像表现"]["主要异常"]["部位"] = line.split("：", 1)[-1]
                elif line.startswith("- 形态："):
                    structured_data["影像表现"]["主要异常"]["形态"] = line.split("：", 1)[-1]
                elif line.startswith("- 密度/信号特征："):
                    structured_data["影像表现"]["主要异常"]["密度"] = line.split("：", 1)[-1]
                elif line.startswith("2. 重要阴性发现"):
                    current_subsection = "阴性发现"
                elif line.startswith("3. 对比既往变化"):
                    current_subsection = "对比变化"
                elif current_subsection == "阴性发现" and line:
                    structured_data["影像表现"]["阴性发现"].append(line)
                elif current_subsection == "对比变化" and line:
                    structured_data["影像表现"]["对比变化"] += line + "\n"
            
            # 初步诊断解析
            elif current_section == "初步诊断":
                if line.startswith("1. 最可能诊断"):
                    current_subsection = "最可能诊断"
                elif line.startswith("2. 鉴别诊断"):
                    current_subsection = "鉴别诊断"
                elif current_subsection == "最可能诊断" and line and line[0].isdigit():
                    structured_data["初步诊断"]["最可能诊断"].append(line.split(". ", 1)[-1])
                elif current_subsection == "鉴别诊断" and line.startswith("-"):
                    structured_data["初步诊断"]["鉴别诊断"].append(line[2:])
            
            # 处理建议解析
            elif current_section == "处理建议":
                if line.startswith("1. 推荐影像检查"):
                    current_subsection = "影像检查"
                elif line.startswith("2. 实验室检查"):
                    current_subsection = "实验室检查"
                elif line.startswith("3. 临床处理方案"):
                    current_subsection = "临床处理"
                elif line.startswith("4. 随访计划"):
                    current_subsection = "随访计划"
                elif current_subsection and line and line[0].isdigit():
                    item = line.split(". ", 1)[-1]
                    structured_data["处理建议"][current_subsection].append(item)

        # ================= 生成报告内容 =================
        # 影像表现部分
        report += "## 影像表现\n"
        if any(structured_data["影像表现"]["主要异常"].values()):
            report += "1. 主要异常：\n"
            report += f"- 部位：{structured_data['影像表现']['主要异常'].get('部位', '未明确')}\n"
            report += f"- 形态：{structured_data['影像表现']['主要异常'].get('形态', '未描述')}\n"
            report += f"- 密度/信号特征：{structured_data['影像表现']['主要异常'].get('密度', '未量化')}\n"
        else:
            report += "1. 主要异常：未发现显著异常\n"
        
        if structured_data["影像表现"]["阴性发现"]:
            report += "2. 重要阴性发现：\n"
            report += "\n".join([f"- {item}" for item in structured_data["影像表现"]["阴性发现"]]) + "\n"
        
        report += f"3. 对比既往变化：{structured_data['影像表现']['对比变化'] or '无可用对比资料'}\n"

        # 初步诊断部分
        report += "\n## 初步诊断\n"
        if structured_data["初步诊断"]["最可能诊断"]:
            report += "1. 最可能诊断（按概率排序）：\n"
            report += "\n".join([f"{i+1}. {d}" for i, d in enumerate(structured_data["初步诊断"]["最可能诊断"])]) + "\n"
        else:
            report += "1. 最可能诊断：需临床进一步确认\n"
        
        if structured_data["初步诊断"]["鉴别诊断"]:
            report += "2. 鉴别诊断及排除依据：\n"
            report += "\n".join([f"- {dd}" for dd in structured_data["初步诊断"]["鉴别诊断"]]) + "\n"

        # 处理建议部分
        report += "\n## 处理建议\n"
        sections = [
            ("1. 推荐影像检查", "影像检查"),
            ("2. 实验室检查", "实验室检查"),
            ("3. 临床处理方案", "临床处理"),
            ("4. 随访计划", "随访计划")
        ]
        for title, key in sections:
            if structured_data["处理建议"][key]:
                report += f"{title}：\n"
                report += "\n".join([f"- {item}" for item in structured_data["处理建议"][key]]) + "\n"

        # ================= 报告元数据 =================
        report += "\n---\n"
        report += f"**报告生成时间**: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
        report += f"**生成系统**: Med-DistillFL医学影像分析系统\n"
        
        return report
    except Exception as e:
        logger.error(f"报告生成失败: {str(e)}")
        return f"报告生成异常，错误信息：{str(e)}\n原始对话历史：{json.dumps(history, ensure_ascii=False)}"


def save_report(history):
    """保存报告文件（修复版）"""
    try:
        report = generate_medical_report(history)
        if not os.path.exists("reports"):
            os.makedirs("reports")
        filename = f"reports/{patient_info['name']}_{int(time.time())}.doc"
        with open(filename, "w", encoding="utf-8") as f:
            f.write(report)
        return filename
    except Exception as e:
        return f"报告保存失败：{str(e)}"
def update_patient_info(name, age, gender, history):
    global patient_info
    patient_info = {
        "name": name or "未填写",
        "age": age or "未知",
        "gender": gender or "未填写",
        "history": history or "无"
    }
    return "患者信息已更新！"

def load_model(args):
    model, _ = AutoModel.from_pretrained(
        args.from_pretrained,
        args=argparse.Namespace(
            fp16=True,
            skip_init=True,
            use_gpu_initialization=torch.cuda.is_available(),
            device='cuda' if torch.cuda.is_available() else 'cpu'
        )
    )
    if args.quant:
        quantize(model.transformer, args.quant)
    return model.eval().cuda() if torch.cuda.is_available() else model.eval()

def main(args):
    global model, tokenizer, current_report
    current_report = ""
    
    model = load_model(args)
    model.add_mixin('auto-regressive', CachedAutoregressiveMixin())
 #   tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)   
    tokenizer = AutoTokenizer.from_pretrained("/data/Matrix/zh/Huggingface-checkpoints/THUDM-chatglm-6b", trust_remote_code=True)
    with gr.Blocks(
        title="Med-DistillFL医学辅助诊断系统",
        css="""
        /* 保持原有CSS样式不变 */
        """
    ) as demo:

        # # ================= 患者信息页面 =================
        # with gr.Column(visible=True, elem_id="patient_info") as info_page:
        #     with gr.Column(elem_classes="content-container"):
        #         gr.Markdown("# 🏥 患者信息录入")
        #         with gr.Row():
        #             with gr.Column():
        #                 name = gr.Textbox(label="患者姓名")
        #                 age = gr.Textbox(label="年龄")
        #                 gender = gr.Radio(["男", "女", "其他"], label="性别")
        #             medical_history = gr.Textbox(label="病史摘要", lines=3)
        #         save_btn = gr.Button("保存信息", variant="primary")
        #         status = gr.Textbox(label="状态", interactive=False)
                
        # ================= 首页 =================
        with gr.Column(visible=True, elem_id="home") as home_page:
            with gr.Column(elem_classes="content-container"):
                gr.Markdown("# 🏥 医学影像智能分析系统 Med-DistillFL")
                # gr.Markdown('# <img src="/data/Matrix/zh/xrayglm/XrayGLM/logo2.ico" width="15" height="15"> 医学影像智能分析系统 Med-DistillFL')
                #gr.Markdown("### 基于多模态大模型的医学影像分析平台")

                # 分类导航
                with gr.Row(elem_classes="category-tags"):
                    gr.Button("全部", elem_classes="tag")
                    gr.Button("放射科", elem_classes="tag")
                    gr.Button("骨科", elem_classes="tag")
                    gr.Button("急诊科", elem_classes="tag")

                # 用例展示
                with gr.Column():
                    gr.Markdown("## 核心分析功能")
                    gr.Markdown("---")

                    # 第一行卡片
                    with gr.Row():
                        with gr.Column(elem_classes="case-card"):
                            gr.Markdown("### 🩺 胸部X光分析")
                            gr.Markdown("""
                            - 肺炎检测与分类
                            - 气胸定量分析
                            - 心脏形态评估
                            - 肋骨骨折定位""")
                            chest_btn = gr.Button("开始分析", variant="primary")

                        with gr.Column(elem_classes="case-card"):
                            gr.Markdown("### 🦴 骨折检测分析")
                            gr.Markdown("""
                            - 骨折类型识别
                            - 愈合阶段评估
                            - 三维重建展示
                            - 治疗方案建议""")
                            fracture_btn = gr.Button("开始分析", variant="primary")

                    # 第二行卡片
                    with gr.Row():
                        with gr.Column(elem_classes="case-card"):
                            gr.Markdown("### 🩻 内窥镜影像分析")
                            gr.Markdown("""
                            - 细胞病理检测
                            - 微生物感染识别
                            - 血管病变评估
                            - 组织病理学检查""")
                            brain_btn = gr.Button("开始分析", variant="primary")

                        with gr.Column(elem_classes="case-card"):
                            gr.Markdown("### 🧠 核磁共振分析")
                            gr.Markdown("""
                            - 脑出血灶检测
                            - 肿瘤定位分析
                            - 脑室系统评估
                            - 卒中区域量化""")
                            abdomen_btn = gr.Button("开始分析", variant="primary")

        # ================= 分析页面 =================
        with gr.Column(visible=False, elem_id="analysis") as analysis_page:
            with gr.Column(elem_classes="content-container"):
                # 导航栏
                with gr.Row():
                    back_btn = gr.Button("← 返回首页", variant="secondary")
                    gr.Markdown("# 🔍 Med-DistillFL影像分析工作台")

                # 患者信息输入区（新增状态组件）
                with gr.Column(visible=True) as patient_info_section:
                    gr.Markdown("## 📝 患者基本信息")
                    with gr.Row():
                        with gr.Column(scale=1):
                            name = gr.Textbox(label="姓名")
                            age = gr.Textbox(label="年龄")
                            gender = gr.Radio(["男", "女"], label="性别")
                        medical_history = gr.Textbox(label="病史摘要", lines=3)
                    save_btn = gr.Button("保存患者信息", variant="primary")
                    status = gr.Textbox(label="保存状态", interactive=False)

                # 主内容区
                with gr.Row():
                    # 输入侧边栏
                    with gr.Column(scale=4):
                        image_prompt = gr.Image(
                            type="filepath",
                            label="上传医学影像",
                            height=400
                        )
                        with gr.Accordion("高级参数设置", open=False):
                            temperature = gr.Slider(0, 1, 0.8, label="生成温度")
                            top_p = gr.Slider(0, 1, 0.4, label="Top P 参数")
                        input_text = gr.Textbox(
                            label="分析指令",
                            placeholder="请输入您的问题...",
                            lines=3
                        )
                        with gr.Row():
                            run_btn = gr.Button("开始分析", variant="primary")
                            clear_btn = gr.Button("清空输入", variant="secondary")

                    # 聊天主界面
                    with gr.Column(scale=6):
                        chatbot = gr.Chatbot(
                            # label="分析对话记录",
                            bubble_full_width=False,
                            height=600
                        )
                
                # 报告生成区
                with gr.Row():
                    report_btn = gr.Button("生成诊疗报告", variant="primary")
                    download_btn = gr.Button("下载报告")
                report_output = gr.Markdown(
                    label="诊疗报告",
                    elem_id="report-output"
                )

        # ================= 事件处理 =================
        save_btn.click(
            update_patient_info,
            [name, age, gender, medical_history],
            status
        )
        
        def reset_to_analysis():
            """重置到分析页面"""
            return [
                gr.update(visible=False),  # home_page
                gr.update(visible=True),   # analysis_page
                [],     # chatbot
                "",     # input_text
                None    # image_prompt
            ]

        def reset_to_home():
            """返回首页并重置状态"""
            return [
                gr.update(visible=True),   # home_page
                gr.update(visible=False),  # analysis_page
                [("系统", "请选择分析功能开始新会话")],  # chatbot
                "",     # input_text
                None    # image_prompt
            ]

        # 绑定功能按钮
        for btn in [chest_btn, fracture_btn, brain_btn, abdomen_btn]:
            btn.click(
                fn=reset_to_analysis,
                outputs=[home_page, analysis_page, chatbot, input_text, image_prompt]
            )

        # 绑定返回按钮
        back_btn.click(
            fn=reset_to_home,
            outputs=[home_page, analysis_page, chatbot, input_text, image_prompt]
        )

        # 绑定分析功能
        run_btn.click(
            request_model,
            [input_text, temperature, top_p, image_prompt, chatbot],
            [input_text, chatbot]
        )
        input_text.submit(
            request_model,
            [input_text, temperature, top_p, image_prompt, chatbot],
            [input_text, chatbot]
        )

        # 绑定清空按钮
        clear_btn.click(
            lambda: ("", [], None),
            outputs=[input_text, chatbot, image_prompt]
        )

        # 新增报告生成事件
        report_btn.click(
            fn=lambda history: generate_medical_report(history),
            inputs=[chatbot],
            outputs=report_output
        )
        
        # 报告下载处理
        def save_report(history):
            report = generate_medical_report(history)
            filename = f"{patient_info['name']}_诊疗报告.doc"
            with open(filename, "w", encoding="utf-8") as f:
                f.write(report)
            return filename
        
        download_btn.click(
            fn=save_report,
            inputs=[chatbot],
            outputs=gr.File(label="下载报告")
        )

    demo.launch(server_name="0.0.0.0", share=args.share)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--quant", type=int, choices=[4,8], default=None)
    parser.add_argument("--share", action="store_true")
    parser.add_argument("--from_pretrained", default="checkpoints")
    args = parser.parse_args()
    main(args)